/*********************************************************
/* mm.c							 *
/* Author: Gabriel Loewen				 *
/* First fit physical memory allocator.			 *
/*********************************************************/

#include <common.h>
#include <mm.h>

block_header_t *blocks;

unsigned char *block_list_addr;

static unsigned int numBlocks;

unsigned int USEDBLOCK;
unsigned int UNUSEDBLOCK;

void * alignAddress(void * ptr);
void startPaging();
unsigned int get_cr0();
void set_cr0(unsigned int value);
unsigned int * createPageDirectory();
unsigned int * createPageTable();

/* Initialize the memory management system.  Assign the location for the block list and create the first
   block header. */
unsigned int * page_table;
unsigned int * page_directory;

void mm_init(unsigned char * base_addr, unsigned char * offset_addr) {

	printf("Starting mm\n");

	UNUSEDBLOCK = (unsigned int)BLOCK_STATUS_UNUSED;
	USEDBLOCK = (unsigned int)BLOCK_STATUS_USED;

	unsigned char *header_offset_addr;
	header_offset_addr = base_addr + (((unsigned int)USABLE_MEM_OFFSET)*sizeof(block_header_t));
	
	memset(base_addr, 0, header_offset_addr);
	printf("Clearing reserved block header address space.\n");

	/* Move the block list to the correct address in memory right after the kernel */
	block_list_addr = (unsigned char *)base_addr;
	blocks = block_list_addr;

	printf("----- Block headers start at 0x%x\n", block_list_addr);

	unsigned char * tmp = block_list_addr + (int)header_offset_addr;

	printf("Usable memory starts at 0x%x\n", tmp);
	printf("Aligning address to 4k boundary.\n");

	unsigned char * startingAddr = alignAddress(tmp);

	printf("Aligned address of usable memory is 0x%x\n", startingAddr);

	blocks[0].base_addr = (unsigned char *)startingAddr;
	blocks[0].offset_addr = offset_addr;
	blocks[0].status = UNUSEDBLOCK;
	numBlocks = 1;

	startPaging();
}

void startPaging() {


	page_directory = createPageDirectory();
	printf("Page directory is at 0x%x\n", page_directory);
	page_table = createPageTable();
	printf("Page table is at 0x%x\n", page_table);
	page_directory[0] = page_table; 
	page_directory[0] |= 3; /* attributes: supervisor level, read/write, present */

	__asm__ volatile("mov %0, %%cr3":: "b"(page_directory));

	/* reads cr0, switches the "paging enable" bit, and writes it back. */
	unsigned int cr0 = get_cr0;
	set_cr0(cr0);

}

unsigned int get_cr0() {
	register unsigned int cr0; 
	__asm__ volatile("mov %%cr0, %0" : "=r" (cr0));
	return(cr0);
}

void set_cr0(unsigned int value) {
	__asm__ volatile("mov %0, %%cr0" : : "r" (value));
}

unsigned int * createPageDirectory() {

	unsigned int * dir = malloc(1024 * sizeof(unsigned int));
	int i = 0;
	for(i = 0; i < 1024; i++) {
		/* attribute: supervisor level, read/write, not present. */
    		dir[i] = 0 | 2; 
	}

	return dir;

}

unsigned int * createPageTable() {

	unsigned int address = 0; 
	unsigned int i;

	unsigned int * table = malloc(1024 * sizeof(unsigned int));
 
	printf("Table: %x\n", table);

	for(i = 0; i < 1024; i++) {
    		table[i] = address | 3; /* attributes: supervisor level, read/write, present. */
    		address = address + (unsigned int)PAGE_FRAME_SIZE; /* advance the address to the next page boundary */
	}

	return table;

}

/* Align the address of ptr to the 4k boundary */	
void * alignAddress(void * ptr) {

	unsigned char * p = ptr;

	while ((unsigned char)p % PAGE_FRAME_SIZE != 0) {
		p++;
	}
	
	return p;
}

/* Returns the difference between 2 numbers */
int diff(unsigned int high, unsigned int low) {
	return high - low;
}

/* Allocate size bytes of memory. */
void * malloc(unsigned int size) {

	int i=0;
	Bool found = FALSE;
	unsigned int retAddress;
	while (i<numBlocks) {
		if (blocks[i].status == UNUSEDBLOCK) {
			printf("Found unused block ar 0x%x\n", blocks[i].base_addr);
			int blockSize = diff(blocks[i].offset_addr,blocks[i].base_addr);
			if (size < blockSize) {
				blocks[i].status = USEDBLOCK;
				unsigned long tmp = blocks[i].base_addr+size+1;
				blocks[numBlocks].base_addr = alignAddress(tmp);
				blocks[numBlocks].offset_addr = blocks[i].offset_addr;
				blocks[numBlocks].status = UNUSEDBLOCK;
				blocks[i].offset_addr = blocks[i].base_addr+size;
				++numBlocks;
				retAddress = blocks[i].base_addr;
				found = TRUE;
				break;
			} else if (blockSize == size) {  /* Fat chance this will ever happen */
				blocks[i].status = USEDBLOCK;
				retAddress = blocks[i].base_addr;
				found = TRUE;
				break;
			} 
		}
		i++;
	}

	if (found == TRUE) {
		return retAddress;
	} else {
		return NULL;
	}
}

/* Allocates n*size bytes and initializes it to 0 */
void * calloc(unsigned int n, unsigned int size) {

	int total = n * size;
	void *p = malloc(total);
	
	if (!p) return NULL;

	return memset(p, 0, total);

}

/* Frees a pointer from memory and attempts to merge neighboring unused blocks */
void free(void * ptr) {

	int i=0;
	while(i<numBlocks) {
		/* We found the address in question */
		if (blocks[i].base_addr == ptr) {
			int size = diff(blocks[i].offset_addr, blocks[i].base_addr);
			blocks[i].status = UNUSEDBLOCK;
			cleanupBlock(i);
			break;
		}
		i++;
	}
}

/* Attempt to merge the address space of neighboring blocks if those blocks are unused */
void cleanupBlock(unsigned int i) {
	
	/* If the block is in the middle of the block list */
	if (i > 0 && i < numBlocks-1) {
		if (blocks[i-1].status == UNUSEDBLOCK &&
	    	    blocks[i+1].status == UNUSEDBLOCK) {
			
			/* Merge blocks[i-1] with blocks[i] and blocks[i+1] 
			   Shift all blocks > blocks[i+1] back by 2 positions*/
			
			blocks[i-1].offset_addr = blocks[i+1].offset_addr;
			shiftBlockHeadersLeft(i+2,2);

		} else if (blocks[i-1].status == USEDBLOCK &&
			   blocks[i+1].status == UNUSEDBLOCK) {
			
			/* Merge blocks[i] with blocks[i+1] 
			   Shift all blocks > blocks[i+1] back by 1 position*/

			blocks[i].offset_addr = blocks[i+1].offset_addr;
			shiftBlockHeadersLeft(i+2,1);

		} else if (blocks[i-1].status == UNUSEDBLOCK &&
			   blocks[i+1].status == USEDBLOCK) {
			
			/* Merge blocks [i-1] with blocks[i]
			   Shift all blocks > blocks[i] back by 1 position*/

			blocks[i-1].offset_addr = blocks[i].offset_addr;
			shiftBlockHeadersLeft(i+1,1);

		}
	} else if (i == 0) { /* If the block is the first block in the list */
		if (blocks[1].status == UNUSEDBLOCK) {
			/* Merge blocks[0] with blocks[1]
			Shift all blocks > blocks[1] back by 1 position*/

			blocks[0].offset_addr = blocks[1].offset_addr;
			shiftBlockHeadersLeft(2,1);

		}
		
	} else if (i == numBlocks-1) { /* If the block is the last block in the list */
		if (blocks[numBlocks-2].status == UNUSEDBLOCK) {
			/* Merge blocks[numBlocks-3] with blocks[numBlocks-2]
			memset blocks[numBlocks-2] with zeroes (Optional)*/
		
			blocks[numBlocks-2].offset_addr = blocks[numBlocks-1].offset_addr;
			--numBlocks;

		}
	}
}

/* Once we merge blocks we want to remove the block headers used by the blocks that are no
   longer in use.  So, we just shift the headers to the left. */
void shiftBlockHeadersLeft(unsigned int blockNum, unsigned int shift) {

	int i = blockNum;
	while (i<numBlocks) {

		blocks[i-shift] = blocks[i];
		++i;

	}
	numBlocks -= shift;

}
